/**
 * 
 */
package edu.umd.clip.lm.model.training.metrics;

import java.util.Arrays;
import java.util.Map;

import org.apache.commons.math.MathException;
import org.apache.commons.math.distribution.*;
import org.w3c.dom.Element;

import edu.umd.clip.lm.model.training.ContextVariable;
import edu.umd.clip.lm.model.training.ContextVariableStats;
import edu.umd.clip.lm.util.ConstCountDistribution;
import edu.umd.clip.lm.util.CountDistribution;
import edu.umd.clip.lm.util.Long2IntMap;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ChiSquaredProbabilityMetrics extends ContextVariableMetrics implements ContextVariableStoppingRule {
	private double maxPValue;
	private final static int MAX_CACHED_DEGREE_OF_FREEDOM = 400;
	private double[] minGoodValues;
	private double[] maxBadValues;
	public static final String P_VAL_ATTR = "p-value";
	
	public ChiSquaredProbabilityMetrics(Element elem) {
		super(elem);
		maxPValue = Double.parseDouble(elem.getAttribute(P_VAL_ATTR));
		init();
	}

	/**
	 * @param orderContextVariables
	 * @param name
	 */
	public ChiSquaredProbabilityMetrics(double maxPValue) {
		super("Chi-squared probability");
		this.maxPValue = maxPValue;
		init();
	}

	private void init() {
		this.minGoodValues = new double[MAX_CACHED_DEGREE_OF_FREEDOM];
		Arrays.fill(minGoodValues, Double.MAX_VALUE);
		
		this.maxBadValues = new double[MAX_CACHED_DEGREE_OF_FREEDOM];
		Arrays.fill(maxBadValues, 0);		
	}
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableMetrics#computeScore(edu.umd.clip.lm.model.training.NewTrainer.ContextVariableStats)
	 */
	@Override
	public double computeScore(ContextVariable ctxVar, ContextVariableStats stat) {
		// compute Chi-squared
		int degreesOfFreedom = (stat.wordCounts.size() - 1) * (stat.contextCounts.size() - 1);
		if (degreesOfFreedom <= 0) {
			return 1.0;
		}
		
		double ChiSquaredSum = computeChiSquaredSum(stat);

		ChiSquaredDistribution chi = new ChiSquaredDistributionImpl(degreesOfFreedom);
		double prob = 1.0;
		try {
			prob = 1.0 - chi.cumulativeProbability(ChiSquaredSum);
		} catch (MathException e1) {
			e1.printStackTrace();
		}
		return prob;
	}

	private double computeProbability(double chiSquaredSum, int degreesOfFreedom) {
		ChiSquaredDistribution chi = new ChiSquaredDistributionImpl(degreesOfFreedom);
		double prob = 1.0;
		try {
			prob = 1.0 - chi.cumulativeProbability(chiSquaredSum);
		} catch (MathException e1) {
			e1.printStackTrace();
		}
		return prob;
	}
	
	private double computeChiSquaredSum(ContextVariableStats stat) {
		double ChiSquaredSum = 0;
		long totalCount = stat.contextCounts.getTotalCount();
		double revTotalCount = 1.0 / totalCount;
		
		long keys[] = stat.x2wDistributions.keys();
		for(long varValue : keys) {
			ConstCountDistribution wordDist = stat.x2wDistributions.get(varValue);
			
			int varExpected = (int) stat.contextCounts.get(varValue);
			
			long seenWordExpectedCount = 0;
			long words[] = wordDist.keys();
			long counts[] = wordDist.values();
			
			for(int i=0; i<words.length; ++i) {
				long word = words[i];
				long wordCount = counts[i];
				
				long wordExpected = stat.wordCounts.get(word);
				
				seenWordExpectedCount += wordExpected;
				
				double expectedCount = revTotalCount * varExpected * wordExpected;
				double diff = wordCount - expectedCount;
				
				double addition = diff * diff / expectedCount;
				if (addition < 0) {
					System.out.println();
				}
				ChiSquaredSum += addition;
			}
			
			// add zero observations
			double zeroObservations = revTotalCount * varExpected * (totalCount - seenWordExpectedCount);
			if (zeroObservations < 0) {
				System.err.println();
			}
			ChiSquaredSum += zeroObservations;
		}
		return ChiSquaredSum;
	}
	
	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableMetrics#needContextVarCounts()
	 */
	@Override
	public boolean needContextVarCounts() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableMetrics#needVarToWordCounts()
	 */
	@Override
	public boolean needVarToWordCounts() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableMetrics#needWordCounts()
	 */
	@Override
	public boolean needWordCounts() {
		return true;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableMetrics#needWordToVarCounts()
	 */
	@Override
	public boolean needWordToVarCounts() {
		return false;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.ContextVariableStoppingRule#isGood(edu.umd.clip.lm.model.training.NewTrainer.ContextVariableStats)
	 */
	@Override
	public boolean isGood(ContextVariable ctxVar, ContextVariableStats stat) {
		int degreesOfFreedom = (stat.wordCounts.size() - 1) * (stat.contextCounts.size() - 1);
		if (degreesOfFreedom <= 0) {
			return false;
		}
		
		double chiSquaredSum = computeChiSquaredSum(stat);
		// check the cache
		if (degreesOfFreedom < MAX_CACHED_DEGREE_OF_FREEDOM) {
			if (chiSquaredSum > minGoodValues[degreesOfFreedom]) {
				return true;
			}
			if (chiSquaredSum < maxBadValues[degreesOfFreedom]) {
				return false;
			}
		}
		
		double pValue = computeProbability(chiSquaredSum, degreesOfFreedom);
		boolean result = pValue <= maxPValue;
		
		// update the cache, no need to worry about concurrent writes since any written value is valid
		if (degreesOfFreedom < MAX_CACHED_DEGREE_OF_FREEDOM) {
			if (result) {
				if (chiSquaredSum < minGoodValues[degreesOfFreedom]) {
					minGoodValues[degreesOfFreedom] = chiSquaredSum;
				}
			} else {
				if (chiSquaredSum > maxBadValues[degreesOfFreedom]) {
					maxBadValues[degreesOfFreedom] = chiSquaredSum;
				}
			}
		}
		
		if (!result) {
			System.err.printf("Chi2 test failed: p=%g, %d vars, %d word, %d total\n", pValue, 
					stat.contextCounts.size(),
					stat.wordCounts.size(),
					stat.wordCounts.getTotalCount());
		}
		return result;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.metrics.ContextVariableMetrics#saveXML(org.w3c.dom.Element)
	 */
	@Override
	public void saveXML(Element elem) {
		super.saveXML(elem);
		elem.setAttribute(P_VAL_ATTR, Double.toString(maxPValue));
	}

	@Override
	public String toString() {
		return "ChiSquaredProbabilityMetrics [maxPValue=" + maxPValue + "]";
	}

}
